import torch
import numpy as np
from agent import DQN
from utils import get_optimal_policy, visual_policy

states = torch.tensor(np.array([[y, x] for y in range(5) for x in range(5)]), dtype=torch.float32)
net = DQN(2, 16, 5)
net.load_state_dict(torch.load("optimal_DQN.pth"))
policy = get_optimal_policy(net, states)
visual_policy(policy)
